# !/usr/bin/env python3
# coding=utf8
"""
程序载入"json配置文件", 程序主要依据"parameters"生成 sqlite 文件,
并将其放在 vnpy.tools.database_cta.CTA_BACKTEST_RESULT_BF_CONT_DBS 目录下,
"""
import argparse
import concurrent.futures
import datetime
import logging
import multiprocessing
import multiprocessing.pool
import os
import pathlib
import queue
import time
import traceback
from typing import Any, Dict, List, Set, Tuple, Type, Optional, Union, Callable, Sequence
from zxt.process_pool_executor_zx import ProcessPoolExecutorZx
import vnpy.tools.database_cta
import vnpy.tools.utility
import vnpy.trader.object
import vnpy.trader.utility
from vnpy.trader.constant import Exchange, Interval
from vnpy.trader.optimize import OUTPUT_FUNC, OptimizationSetting
from vnpy_ctastrategy.backtesting import BacktestingEngine
from vnpy_zxtools.ctabacktester.cta_config import CtaConfigLite, CtaConfig


def get_cta_strategy_pickle_data(engine: BacktestingEngine) -> Optional[bytes]:
    """"""
    bin_data: Optional[bytes] = None

    cta_dump_data: vnpy.tools.database_cta.CtaDumpData = vnpy.tools.database_cta.CtaDumpData(
        strategy_name=engine.strategy_class.__name__,
        interval=engine.interval,
        exchange=engine.exchange,
        symbol=engine.symbol,
        setting=engine.strategy.get_parameters(),
    )
    sqlite_filepath: pathlib.Path = vnpy.tools.database_cta.CtaDumpDataTool.CtaDumpData_2_filepath(cta_dump_data=cta_dump_data)
    if sqlite_filepath.is_file():
        database: vnpy.tools.database_cta.CtaBacktestDatabase = vnpy.tools.database_cta.get_database()
        database.reinitialize(sqlite_filepath=sqlite_filepath)

        lines: List[vnpy.tools.database_cta.CtaStrategyPickleData] = database.load_cta_strategy_pickle_data(
            strategy_name=engine.strategy_class.__name__,
            interval=engine.interval,
            exchange=engine.exchange,
            symbol=engine.symbol,
            setting=engine.strategy.get_parameters(),
        )
        assert len(lines) <= 1

        bin_data: Optional[bytes] = lines[0].data if lines else None

    return bin_data


class Consumer_deprecated(object):
    """"""

    def __init__(self, queue_is_none: bool = False) -> None:
        self.queue: Optional[queue.Queue] = None if queue_is_none else multiprocessing.Manager().Queue()
        self.process: multiprocessing.Process = None

    def start(self) -> None:
        """"""
        if (self.queue is not None) and (self.process is None):
            self.process = multiprocessing.Process(target=self.process_entry, args=(self.queue,))
            self.process.start()

    def join(self) -> None:
        """"""
        if self.process is not None:
            self.process.join()

    def put_none(self) -> None:
        """"""
        if self.queue is not None:
            self.queue.put(item=None)

    @classmethod
    def process_entry(cls, que: queue.Queue) -> None:
        """要求该函数能自动结束"""
        while cls.batch(que=que):
            time.sleep(0.1)

    @classmethod
    def batch(cls, que: queue.Queue) -> bool:
        """"""
        is_alive: bool = True

        items: List[vnpy.tools.database_cta.CtaDumpData] = []

        while True:
            try:
                item = que.get_nowait()

                if item is None:  # 一旦收到None就认为工作完成了,
                    is_alive: bool = False
                    break

                assert isinstance(item, vnpy.tools.database_cta.CtaDumpData)
                items.append(item)

            except queue.Empty as ex:
                break

        if items:
            vnpy.tools.database_cta.CtaDumpDataTool.save_to_database_single(items=items)

        return is_alive


def initialize_logging():
    """"""
    c_handler: logging.StreamHandler = logging.StreamHandler()
    f_handler: logging.FileHandler = logging.FileHandler(filename=f"{__file__}.log", mode="a", encoding="utf8")
    logging.basicConfig(level=logging.DEBUG, format="%(asctime)s %(levelname)s - %(message)s", handlers=[c_handler, f_handler])
    logging.getLogger("peewee").setLevel(logging.INFO)


def logging_log(msg) -> None:
    """"""
    logging.info(msg)


def evaluate_zx(
    old_engine: BacktestingEngine,
    setting: dict,
) -> None:
    """
    修改自 vnpy_ctastrategy/vnpy_ctastrategy/backtesting.py 的 evaluate 函数,
    """
    # This function does nothing if the root logger already has handlers configured, unless ...
    logging.basicConfig(level=logging.DEBUG, format="%(message)s")
    logging.getLogger("peewee").setLevel(logging.INFO)

    engine: BacktestingEngine = BacktestingEngine()

    # TODO: 为什么 print 无法输出到控制台, logging 却可以呢?
    engine.output = lambda msg: logging.info(f"{datetime.datetime.now()}{engine.pid:>9d}\t{msg}")

    engine.set_parameters(
        vt_symbol=old_engine.vt_symbol,
        interval=old_engine.interval,
        start=old_engine.start,
        rate=old_engine.rate,
        slippage=old_engine.slippage,
        size=old_engine.size,
        pricetick=old_engine.pricetick,
        capital=old_engine.capital,
        end=old_engine.end,
        mode=old_engine.mode,
        risk_free=old_engine.risk_free,
        annual_days=old_engine.annual_days,
        do_continue=old_engine.do_continue,
    )

    engine.add_strategy(old_engine.strategy_class, setting)
    bin_data: Optional[bytes] = get_cta_strategy_pickle_data(engine=engine)
    engine.load_strategy(bin_data=bin_data)
    if engine.datetime is not None:
        # 首次执行时, 刚好指定了春节期间的时间段, 因为这个时间段没有数据, 所以没有进入回测逻辑, 所以 BacktestingEngine.datetime 尚未赋值,
        # 回测结束后, 紧接着就将 bin_data 落库了, 再次载入 bin_data 后, 就会出现 BacktestingEngine.datetime 为 None 的情况,
        assert bin_data is not None
        engine.start, engine.end = engine.fix_start_end(start=engine.datetime.replace(tzinfo=None), end=engine.end)
    engine.output(f"evaluate_zx, start={engine.start}, end={engine.end}, setting={engine.strategy.get_parameters()}")
    engine.load_data()
    engine.run_backtesting()
    engine.calculate_result()

    cta_dump_data: vnpy.tools.database_cta.CtaDumpData = vnpy.tools.database_cta.CtaDumpDataTool.generate_cta_dump_data(
        engine=engine,
        is_cta_trade_data_list=False,
        is_cta_daily_result_data_list=True,
        is_cta_statistics_data=False,
        is_cta_strategy_pickle_data=True,
    )

    vnpy.tools.database_cta.CtaDumpDataTool.save_to_database_multi(cta_dump_data=cta_dump_data)

    engine.output(f"evaluate_zx, start={engine.start}, end={engine.end}, setting={engine.strategy.get_parameters()}, done")

    return None


def run_bf_optimization_v2(
    engine: BacktestingEngine,
    optimization_setting: OptimizationSetting,
    max_workers: int = None,
    output: OUTPUT_FUNC = print
) -> None:
    """
    修改自 vnpy/vnpy/trader/optimize.py 的 run_bf_optimization_v2 函数,
    """
    """Run brutal force optimization"""
    # 假定程序续跑前后 batch_size 无变化, 这一批 setting 如果都执行完了, 那是真的都执行完了,
    # 这一批 setting 如果执行了一部分的时候, 程序挂了, 那这一批 setting 都重新执行就行了,
    # 让对象 optimization_setting 从外面带进来一个额外的"backtest_resume"字段,
    if hasattr(optimization_setting, "backtest_resume"):
        backtest_resume: int = optimization_setting.backtest_resume
    else:
        backtest_resume: int = 0

    settings_md5: str = optimization_setting.checksum()
    output(f"settings_md5={settings_md5}")  # 如果续跑, 对比它可知参数是否发生变化,

    len_settings: int = optimization_setting.settings_num()

    output("开始执行穷举算法优化")
    output(f"参数优化空间：{len_settings}")

    start: int = time.perf_counter()

    results_num: int = 0

    percentage: float = results_num / len_settings * 100
    output("穷举算法优化进度：{:d}/{:d}, {:.2f}%".format(results_num, len_settings, percentage))

    batch_size: int = 10_000
    for _settings in optimization_setting.generate_settings_yield(batch_size=batch_size):
        if int(results_num / batch_size) < int(backtest_resume / batch_size):
            results_num += len(_settings)
            output("穷举算法优化进度：{:d}/{:d}, backtest_resume={:d}".format(results_num, len_settings, backtest_resume))
            continue

        if False:
            for setting in _settings:
                evaluate_zx(engine, setting)
            break

        with concurrent.futures.ProcessPoolExecutor(
            max_workers=max_workers,
            mp_context=multiprocessing.get_context("spawn")
        ) as executor:

            fs = {executor.submit(evaluate_zx, engine, setting) for setting in _settings}

            for future in concurrent.futures.as_completed(fs):
                try:
                    future.result()  # 如果任务异常了, 会通过它感知到,
                except Exception as exception:
                    output(f"exception={exception}; " + traceback.format_exc())

                results_num += 1

                if results_num % 1000 == 0:
                    percentage: float = results_num / len_settings * 100
                    output("穷举算法优化进度：{:d}/{:d}, {:.2f}%".format(results_num, len_settings, percentage))

    percentage: float = results_num / len_settings * 100
    output("穷举算法优化进度：{:d}/{:d}, {:.2f}%".format(results_num, len_settings, percentage))

    end: int = time.perf_counter()
    cost: int = int((end - start))
    output(f"穷举算法优化完成，耗时{cost}秒")

    return None


def run_bf_optimization(cta_config: CtaConfig):
    """
    多进程优化, 基于 vnpy/examples/cta_backtesting/backtesting_demo.ipynb 进行修改,
    遗传算法(Genetic Algorithm, GA), 蛮力优化(brutal force optimization),
    """
    assert (0 <= cta_config.start.weekday() <= 4) and (0 <= cta_config.end.weekday() <= 4)

    engine: BacktestingEngine = BacktestingEngine()

    engine.output = logging_log

    engine.set_parameters(
        vt_symbol=cta_config.vt_symbol,
        interval=cta_config.interval,
        start=cta_config.start,
        rate=cta_config.rate,
        slippage=cta_config.slippage,
        size=cta_config.size,
        pricetick=cta_config.pricetick,
        capital=cta_config.capital,
        end=cta_config.end,
        mode=cta_config.mode,
        risk_free=cta_config.risk_free,
        annual_days=cta_config.annual_days,
        do_continue=cta_config.do_continue,
    )

    engine.add_strategy(cta_config.strategy_class, {})

    optimization_setting: OptimizationSetting = OptimizationSetting()

    if isinstance(cta_config.backtest_resume, int) and 0 < cta_config.backtest_resume:
        optimization_setting.backtest_resume = cta_config.backtest_resume

    for name, d in cta_config.parameters.items():
        start, last, step = d["start"], d["last"], d["step"]
        if start == last:
            step = last = None
        optimization_setting.add_parameter(name=name, start=start, end=last, step=step)

    if cta_config.strategy_class.get_class_parameters().get(Interval.__name__.lower(), None) in Interval._value2member_map_:
        # 本来应该在 cta_setting.parameters 里面配置 Interval.__name__.lower() 的, 但是已经有 cta_setting.interval 了,
        # 为了避免重复配置, 所以在这里自动猜测(参数里面是否有 interval), 如果有, 就自动取用 cta_setting.interval 的值,
        optimization_setting.add_parameter(name=Interval.__name__.lower(), start=cta_config.interval.value, end=None, step=None)

    run_bf_optimization_v2(engine=engine, optimization_setting=optimization_setting, output=logging_log)


def run_bf_optimization_test(cta_config: CtaConfig):
    """"""
    assert cta_config.start == datetime.datetime.strptime(cta_config.start.strftime("%Y-%m-%d"), "%Y-%m-%d")
    assert cta_config.end == datetime.datetime.strptime(cta_config.end.strftime("%Y-%m-%d"), "%Y-%m-%d")
    #
    assert (0 <= cta_config.start.weekday() <= 4) and (0 <= cta_config.end.weekday() <= 4)
    #
    curr: datetime.datetime = cta_config.start
    last: datetime.datetime = cta_config.end
    while curr <= last:
        cta_config.end = curr
        run_bf_optimization(cta_config=cta_config)
        curr += datetime.timedelta(days={4: 3, 5: 2}.get(curr.weekday(), 1))  # 寻找下一个工作日


# 查看 pickle.dumps 生成的二进制数据的MD5:
# sqlite3 -csv -newline '#' database_cta_cont.db "SELECT md5 FROM dbctastrategypickledata ORDER BY setting ASC LIMIT 1;" | md5sum
# find . -type f -name "*.db" -exec sqlite3 -csv {} "SELECT setting, md5 FROM dbctastrategypickledata;" \; | sort | md5sum
# 测试"续跑模式"的函数: run_bf_optimization_test,
if __name__ == "__main__":
    initialize_logging()

    logging.info(f"++++++++++ START ++++++++++")
    logging.info(f"pid_tag, getpid={os.getpid()}, getppid={os.getppid()}")

    parser = argparse.ArgumentParser()
    parser.add_argument("--file", "-f", type=str, default="", help="json配置文件")
    args = parser.parse_args()

    json_file: pathlib.Path = pathlib.Path(args.file)

    if not json_file.is_file():
        logging.info(f"json_file={json_file.absolute()}, not is_file, will exit...")
        exit(1)

    logging.info(f"json_file={json_file.absolute()},")

    cta_config_lite: CtaConfigLite = CtaConfigLite.load_from_file(filepath=json_file)
    logging.info(cta_config_lite)

    cta_config: CtaConfig = cta_config_lite.to_cta_config()
    logging.info(cta_config)

    assert (0 <= cta_config.start.weekday() <= 4) and (0 <= cta_config.end.weekday() <= 4)

    run_bf_optimization(
        cta_config=cta_config,
    )
    # 请在当前文件搜索 vnpy.tools.database_cta.CtaDumpDataTool.generate_cta_dump_data 然后择需打开数据的开关,
